import math as pm
import numpy as np
from numpy import genfromtxt
from scipy import stats
import matplotlib.pyplot as plt
import mpl_toolkits.mplot3d.axes3d as p3
import h5py
import mpi4py

PREFIX_U = 'u_r_'
PREFIX_V = 'v_r_'
PREFIX_RHO = 'rho_r_'
PATH = '/home/behollis/DATA/in/ts00050/'
DPATH = '/home/behollis/DATA/out/lockExSt/ts00050/'

gUv = dict()   
DIM = 2

def lerp(a, b, w):
    return a + (b - a) * w

def readMember(mem):
    ''' Read numpy matrix from disk and return tuple of realization data. '''
    filename_u = PREFIX_U + str(mem) + '.txt'
    filename_v = PREFIX_V + str(mem) + '.txt'
    #filename_rho = PREFIX_RHO + str(mem) + '.txt'
    
    u = genfromtxt(PATH + filename_u)
    u = -1.0*u[0:127,0:127]
    
    v = genfromtxt(PATH + filename_v)
    #rho = genfromtxt(PATH + filename_rho)
    
    return (u, v)# , rho)

def getCoordParts(ppos=[0.0,0.0]):
    #decompose fract / whole from particle position
    ppos_parts = [[0.0,0.0],[0.0,0.0]] #[fract,whole] for each x,y comp
    ppos_parts[0][0] = pm.modf(ppos[0])[0];ppos_parts[0][1] = pm.modf(ppos[0])[1]
    ppos_parts[1][0] = pm.modf(ppos[1])[0];ppos_parts[1][1] = pm.modf(ppos[1])[1]
    
    return ppos_parts 

def getGridPoints(ppos=[0.0,0.0]):
    #assume grid points are defined by integer indices
    
    ppos_parts = getCoordParts(ppos)
    
    #print "quantile alpha x: " + str( ppos_parts[0][0] )
    #print "quantile alpha y: " + str( ppos_parts[1][0] )
    
    # grid point numbers:
    #
    # (2)---(3)
    # |      |
    # |      |
    # (0)---(1)
    
    #find four corner grid point indices, numbered from gpt0 = (bottom, left) TO gpt3 = (top, right)
    #calculated from whole parts 
    gpt0 = [ppos_parts[0][1], ppos_parts[1][1]]
    gpt1 = [ppos_parts[0][1] + 1, ppos_parts[1][1]]
    gpt2 = [ppos_parts[0][1], ppos_parts[1][1] + 1]
    gpt3 = [ppos_parts[0][1] + 1, ppos_parts[1][1] + 1]
    
    return gpt0, gpt1, gpt2, gpt3

def getVclinSamples(gpt0, gpt1, gpt2, gpt3):
    gpt0_dist = np.zeros(shape=(DIM,MEM))
    gpt1_dist = np.zeros(shape=(DIM,MEM))
    gpt2_dist = np.zeros(shape=(DIM,MEM))
    gpt3_dist = np.zeros(shape=(DIM,MEM))
    
    for idx in range(1,MEM):
        gpt0_dist[0][idx] = gUv[idx][0][gpt0[0]][gpt0[1]]
        gpt0_dist[1][idx] = gUv[idx][1][gpt0[0]][gpt0[1]]
        
        gpt1_dist[0][idx] = gUv[idx][0][gpt1[0]][gpt1[1]]
        gpt1_dist[1][idx] = gUv[idx][1][gpt1[0]][gpt1[1]]
        
        gpt2_dist[0][idx] = gUv[idx][0][gpt2[0]][gpt2[1]]
        gpt2_dist[1][idx] = gUv[idx][1][gpt2[0]][gpt2[1]]
        
        gpt3_dist[0][idx] = gUv[idx][0][gpt3[0]][gpt3[1]]
        gpt3_dist[1][idx] = gUv[idx][1][gpt3[0]][gpt3[1]]
   
    return gpt0_dist, gpt1_dist, gpt2_dist, gpt3_dist


def interpVelFromEnsemble(ppos=[0.0,0.0]):
    
    ppos_parts = getCoordParts(ppos)
    
    #find grid cell that contains position to be interpolated
    gpt0, gpt1, gpt2, gpt3 = getGridPoints(ppos)
    
    gpt0_dist, gpt1_dist, gpt2_dist, gpt3_dist = getVclinSamples(gpt0, gpt1, gpt2, gpt3)
    
    #lerp ensemble samples
    lerp_u_gp0_gp1 = lerp( np.asarray( gpt0_dist[0] ), np.asarray( gpt1_dist[0]), w = ppos_parts[0][0] )
    lerp_u_gp2_gp3 = lerp( np.asarray( gpt2_dist[0] ), np.asarray( gpt3_dist[0]), w = ppos_parts[0][0] ) 
    lerp_u = lerp( np.asarray(lerp_u_gp0_gp1), np.asarray(lerp_u_gp2_gp3), w = ppos_parts[1][0] )  
    
    lerp_v_gp0_gp1 = lerp( np.asarray(gpt0_dist[1] ), np.asarray(gpt1_dist[1]), w = ppos_parts[0][0] )
    lerp_v_gp2_gp3 = lerp( np.asarray(gpt2_dist[1] ), np.asarray(gpt3_dist[1]), w = ppos_parts[0][0] ) 
    lerp_v = lerp( np.asarray(lerp_v_gp0_gp1), np.asarray(lerp_v_gp2_gp3), w = ppos_parts[1][0] )  
        
    k = stats.kde.gaussian_kde((lerp_u,lerp_v)) 
    return k, lerp_u.min(axis=0), lerp_u.max(axis=0), lerp_v.min(axis=0), lerp_v.max(axis=0)

def plotDistro(kde,x_min_max,y_min_max,title=''):
    
    div = 100j
    div_real = 100
    
    # Regular grid to evaluate kde upon
    x_flat = np.r_[x_min_max[0]:x_min_max[1]:div]
    y_flat = np.r_[y_min_max[0]:y_min_max[1]:div]
    x,y = np.meshgrid(x_flat,y_flat)
    
    #positions = np.vstack([x_flat.ravel(), y_flat.ravel()])
    
    grid_coords = np.append(x.reshape(-1,1),y.reshape(-1,1),axis=1)
    z = kde(grid_coords.T)
    
    z = z.reshape(div_real,div_real)
    
    fig = plt.figure()
                        
    ax = fig.gca(projection='3d')
    surf = ax.plot_surface(x, y, z, rstride=2, cstride=2, alpha=0.1, linewidth=0, antialiased=True, color='b')#,cmap=cm.spectral)
    
    ax.set_xlabel('u')
    ax.set_ylabel('v')    
    ax.set_zlabel('density')
    
    plt.show() 
    
def readStreamlines(x, y, f):
    slines = list()
    
    groups = f.keys()
    
    for mem in groups:
        dir = str(mem).zfill(4) + '/x' + str(x).zfill(3) + '/y' + str(y).zfill(3)
    
        xlst = list(f[dir][0])
        ylst = list(f[dir][1])
        zlst = 0.0 * len(list(f[dir][0]))
    
        slines.append([xlst,ylst,zlst])
        
    return slines   

import string

from OpenGL.GL import *
from OpenGL.GLUT import *
from OpenGL.GLU import *
import sys
from Image import *

# Some api in the chain is translating the keystrokes to this octal string
# so instead of saying: ESCAPE = 27, we use the following.
ESCAPE = '\033'

# Number of the glut window.
window = 0

# Rotations for cube. 
xrot=yrot=zrot=0.0
xspeed=yspeed=zspeed=0.0

texture = 0

object=0

def LoadTextures():
    #global texture
    image = open('../img/pfield.png')
    
    ix = image.size[0] 
    iy = image.size[1] 
    image = image.tostring("raw", "RGB", 0, -1)
    
    
    
    # Create Texture    
    glBindTexture(GL_TEXTURE_2D, glGenTextures(1))   # 2d texture (x and y size)
    
    glPixelStorei(GL_UNPACK_ALIGNMENT,1)
    glTexImage2D(GL_TEXTURE_2D, 0, 3, ix, iy, 0, GL_RGB, GL_UNSIGNED_BYTE, image)
    glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP)
    glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP)
    glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT)
    glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT)
    glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
    glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)
    glTexEnvf(GL_TEXTURE_ENV, GL_TEXTURE_ENV_MODE, GL_DECAL)
    

    
# The main drawing function. 
def DrawGLScene():
    global texture, prob, slines
    
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)    # Clear The Screen And The Depth Buffer
    
    glLoadIdentity()
    gluLookAt(0,0,-10, 125.0 / 2.0, 125.0 / 2.0, 1, 0, 1, 0)
      
    
    
     
    '''
    glBegin(GL_QUADS)
    
    #glNormal2f( 0.0, 0.0, 1.0)

    glTexCoord2f(0.0, 0.0)
    glVertex2f(-W, -H, 1.0)
    
    glTexCoord2f(1.0, 0.0)
    glVertex2f( W, -H, 1.0)
    
    glTexCoord2f(1.0, 1.0)
    glVertex2f( W, H, 1.0)
    
    glTexCoord2f(0.0, 1.0)
    glVertex2f(-W, H, 1.0)

    glEnd()
    '''

    #for slidx, sl in enumerate(slines):
    for slidx in range(0, 2):
        sl = slines[slidx]
        
        glBegin(GL_LINE_STRIP)
        glColor4f(1.0, 1.0, 1.0, 1.0)
        for ptidx in range(0, len(sl[0])):
            
            pt = [ sl[0][ptidx], sl[1][ptidx] ]
            glVertex3f(pt[0], pt[1], 1.0)
            #glVertex2f(ptidx, 0.0)
            #print pt[0]
            #print pt[1]
            #print len(sl[0])
            
        glEnd()
        
    
        
    
    #  since this is double buffered, swap the buffers to display what just got drawn. 
    glutSwapBuffers()


#A general OpenGL initialization function.  Sets all of the initial parameters. 
def InitGL(Width, Height):                # We call this right after our OpenGL window is created.
    global quadratic
    LoadTextures()
     
    glEnable(GL_TEXTURE_2D)
    glClearColor(0.0, 0.0, 0.0, 0.0)    # This Will Clear The Background Color To Black
    glClearDepth(1.0)                    # Enables Clearing Of The Depth Buffer
    glDepthFunc(GL_LESS)                # The Type Of Depth Test To Do
    glEnable(GL_DEPTH_TEST)                # Enables Depth Testing
    glShadeModel(GL_SMOOTH)                # Enables Smooth Color Shading

##    glPolygonMode(GL_FRONT, GL_LINE)
##    glPolygonMode(GL_BACK, GL_LINE)
    
    glTexGeni(GL_S, GL_TEXTURE_GEN_MODE, GL_SPHERE_MAP)
    glTexGeni(GL_T, GL_TEXTURE_GEN_MODE, GL_SPHERE_MAP)
    
    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()                    # Reset The Projection Matrix
                                        # Calculate The Aspect Ratio Of The Window
    #gluPerspective(45.0, float(Width)/float(Height), 0.1, 100.0)
    
    W = 125
    H = 125
    gluOrtho2D(0, W, 0, H )
    
    

    glMatrixMode(GL_MODELVIEW)

# The function called whenever a key is pressed. Note the use of Python tuples to pass in: (key, x, y)  
def keyPressed(key,x,y):
    global object
    # If escape is pressed, kill everything.
    key = string.upper(key)
    if key == ESCAPE:
        sys.exit()
    '''
    elif key == 'L':
        light = not light
    elif key == 'T': #  switch the texture
        texture_num = (texture_num + 1) % 3
    elif key == 'O': #  switch the object
        object = (object + 1) % 7
    '''
        
# The function called when our window is resized (which shouldn't happen if you enable fullscreen, below)
def ReSizeGLScene(Width, Height):
    if Height == 0:                        # Prevent A Divide By Zero If The Window Is Too Small 
        Height = 1

    glViewport(0, 0, Width, Height)        # Reset The Current Viewport And Perspective Transformation
    glMatrixMode(GL_PROJECTION)
    
    glLoadIdentity()
    
    W = 125
    H = 125
    #gluOrtho2D(0, W, 0, H )
    gluPerspective(45.0, float(Width)/float(Height), 0.1, 100.0)
    glMatrixMode(GL_MODELVIEW)
        
def main():
    global window
    glutInit(sys.argv)

    # Select type of Display mode:   
    #  Double buffer 
    #  RGBA color
    # Alpha components supported 
    # Depth buffer
    glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_DEPTH)
    
    # get a 640 x 480 window 
    glutInitWindowSize(500, 500)
    
    # the window starts at the upper left corner of the screen 
    glutInitWindowPosition(500, 500)
    
    # Okay, like the C version we retain the window id to use when closing, but for those of you new
    # to Python (like myself), remember this assignment would make the variable local and not global
    # if it weren't for the global declaration at the start of main.
    window = glutCreateWindow("Research Paper: streamline liklihood")

    # Register the drawing function with glut, BUT in Python land, at least using PyOpenGL, we need to
    # set the function pointer and invoke a function to actually register the callback, otherwise it
    # would be very much like the C version of the code.    
    glutDisplayFunc(DrawGLScene)
    
    # Uncomment this line to get full screen.
    #glutFullScreen()

    # When we are doing nothing, redraw the scene.
    glutIdleFunc(DrawGLScene)
    
    # Register the function called when our window is resized.
    glutReshapeFunc(ReSizeGLScene)
    
    # Register the function called when the keyboard is pressed.  
    glutKeyboardFunc(keyPressed)

    # Initialize our window. 
    InitGL(500, 500)
    
    

    # Start Event Processing Engine  
    glutMainLoop()
    
def calcStreamlineProbs(seed_pt):
    slines_prob = dict()
    
    for m in range(1, MEM):
        gUv[m] = readMember(m) 
    
    
#    for slidx, sl in enumerate(slines):
    for slidx in range(0, 2):
        sl = slines[slidx]
        tot = 0
        slines_prob[slidx] = [ None, [] ]
        for ptidx in range(0, len(sl[0])):
            
            pt = [ sl[0][ptidx], sl[1][ptidx] ]
            
            kde, min_u, max_u, min_v, max_v = interpVelFromEnsemble(pt)
            #plotDistro( kde, (min_u,max_u), (min_v,max_v), '') 
            #band_width = MEM**(-1./(DIM+4))
            epsx = 0.005 * pm.fabs(max_u - min_u)
            epsy = 0.005 * pm.fabs(max_v - min_v)
              
            m = slidx + 1# member and slidx are synonymous
            
            u = gUv[m][0][pt[0],pt[1]]
            v = gUv[m][1][pt[0],pt[1]]
            low_bounds = [u - epsx, v - epsy]
            high_bounds = [u + epsx, v + epsy]
            prob = kde.integrate_box( low_bounds, high_bounds )
            tot += prob
            slines_prob[slidx][1].append(prob)
                
        print 'streamline {0}: {1}'.format(slidx, tot)
        slines_prob[slidx][0] = tot
        
        
    return slines_prob, slines
    
prob = None
slines = None
MEM = 200
if __name__ == "__main__":
    
    global prob, slines
    
    f = h5py.File(DPATH+'0.0.hdf5', 'r')
    
    seed_pt = [40, 25]
    slines = readStreamlines(seed_pt[0], seed_pt[1], f)
    

    #prob = calcStreamlineProbs(seed_pt)
        
    print "Hit ESC key to quit."
    main()
    
    
    

            